package utils

import (
	"sync"
	"time"
)

// ConnectionHealth 连接健康状态
type ConnectionHealth struct {
	LastHeartbeat time.Time
	IsHealthy     bool
	FailureCount  int
}

// HealthManager 健康管理器
type HealthManager struct {
	connections map[string]*ConnectionHealth
	mutex       sync.RWMutex
	timeout     time.Duration
	maxFailures int
}

// NewHealthManager 创建健康管理器
func NewHealthManager(timeout time.Duration, maxFailures int) *HealthManager {
	return &HealthManager{
		connections: make(map[string]*ConnectionHealth),
		timeout:     timeout,
		maxFailures: maxFailures,
	}
}

// UpdateHeartbeat 更新心跳时间
func (hm *HealthManager) UpdateHeartbeat(connID string) {
	hm.mutex.Lock()
	defer hm.mutex.Unlock()

	if health, exists := hm.connections[connID]; exists {
		health.LastHeartbeat = time.Now()
		health.IsHealthy = true
		health.FailureCount = 0
	} else {
		hm.connections[connID] = &ConnectionHealth{
			LastHeartbeat: time.Now(),
			IsHealthy:     true,
			FailureCount:  0,
		}
	}
}

// MarkUnhealthy 标记连接不健康
func (hm *HealthManager) MarkUnhealthy(connID string) {
	hm.mutex.Lock()
	defer hm.mutex.Unlock()

	if health, exists := hm.connections[connID]; exists {
		health.IsHealthy = false
		health.FailureCount++
	} else {
		hm.connections[connID] = &ConnectionHealth{
			LastHeartbeat: time.Now(),
			IsHealthy:     false,
			FailureCount:  1,
		}
	}
}

// CheckHealth 检查连接健康状态
func (hm *HealthManager) CheckHealth(connID string) bool {
	hm.mutex.RLock()
	defer hm.mutex.RUnlock()

	health, exists := hm.connections[connID]
	if !exists {
		return false
	}

	// 检查是否超时
	if time.Since(health.LastHeartbeat) > hm.timeout {
		return false
	}

	// 检查失败次数
	if health.FailureCount >= hm.maxFailures {
		return false
	}

	return health.IsHealthy
}

// RemoveConnection 移除连接
func (hm *HealthManager) RemoveConnection(connID string) {
	hm.mutex.Lock()
	defer hm.mutex.Unlock()
	delete(hm.connections, connID)
}

// GetUnhealthyConnections 获取不健康的连接
func (hm *HealthManager) GetUnhealthyConnections() []string {
	hm.mutex.RLock()
	defer hm.mutex.RUnlock()

	var unhealthy []string
	now := time.Now()

	for connID, health := range hm.connections {
		if !health.IsHealthy || 
		   now.Sub(health.LastHeartbeat) > hm.timeout || 
		   health.FailureCount >= hm.maxFailures {
			unhealthy = append(unhealthy, connID)
		}
	}

	return unhealthy
}

// CleanupStaleConnections 清理过期的连接
func (hm *HealthManager) CleanupStaleConnections(maxAge time.Duration) int {
	hm.mutex.Lock()
	defer hm.mutex.Unlock()

	now := time.Now()
	staleCount := 0

	for connID, health := range hm.connections {
		if now.Sub(health.LastHeartbeat) > maxAge {
			delete(hm.connections, connID)
			staleCount++
		}
	}

	return staleCount
}